#include "is_real_obj_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_is_real_obj_custom.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void is_real_obj_custom(GM_ADDR x, GM_ADDR out, IsRealObjCustomTilingData tiling);
#endif
extern void GenerateTilingData(uint8_t *tilingBuf, uint32_t blockDim);

int32_t main(int32_t argc, char *argv[])
{
    constexpr uint32_t BLOCK_DIM = 8;
    constexpr uint32_t DATA_TYPE_SIZE[] = {2, 2, 4, 1, 2, 4, 1, 2, 4, 8, 8, 1, 8, 0, 1, 1, 8, 16, 1, 2,
     4, 1, 2, 0, 0, 0, 0, 0, 1, 1, 1, 1, 4, 1, 1, 1, 1, 1, 1, 1, 1, 0};
    uint8_t *tiling = nullptr;
    size_t tilingSize = 17 * sizeof(uint32_t);

#ifdef ASCENDC_CPU_DEBUG
    tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);

#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost, *outHost;
    uint8_t *xDevice, *outDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#endif
    GenerateTilingData(tiling, BLOCK_DIM);
    uint32_t dataTypeSize = DATA_TYPE_SIZE[reinterpret_cast<IsRealObjCustomTilingData *>(tiling)->dataType];
    size_t inputByteSize = reinterpret_cast<IsRealObjCustomTilingData *>(tiling)->totalLength * dataTypeSize;
    size_t outputByteSize = sizeof(int8_t);  // 只输出一个布尔值

#ifdef ASCENDC_CPU_DEBUG
    constexpr uint32_t DATA_BLOCK_BYTE = 32;

    uint8_t *x = (uint8_t *)AscendC::GmAlloc((inputByteSize + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);
    uint8_t *out = (uint8_t *)AscendC::GmAlloc((outputByteSize + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    
    ICPU_RUN_KF(is_real_obj_custom, BLOCK_DIM, x, out,
                *reinterpret_cast<IsRealObjCustomTilingData *>(tiling));

    WriteFile("./output/output_out.bin", out, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)out);
    AscendC::GmFree((void *)tiling);

#else
    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&outHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&outDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(is_real_obj_custom)(BLOCK_DIM, stream, xDevice, outDevice,
                                       reinterpret_cast<IsRealObjCustomTilingData *>(tiling));
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(outHost, outputByteSize, outDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_out.bin", outHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(outDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(outHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}